import torch
import torch.nn as nn
from torch.distributions import Normal

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class ActorNetwork(nn.Module):
    """独立Actor网络，处理局部观测"""
    def __init__(self, state_dim, action_dim, hidden_dim=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, action_dim),
            nn.Tanh()  # 输出范围[-1,1]
        )
        self.action_scale = None
        self.action_bias = None

    def set_action_scale(self, low, high):
        self.action_scale = torch.tensor((high - low)/2, dtype=torch.float32, device=device)
        self.action_bias = torch.tensor((high + low)/2, dtype=torch.float32, device=device)

    def forward(self, state):
        raw_action = self.net(state)
        return raw_action * self.action_scale + self.action_bias  # 缩放到实际动作范围

class CentralizedCritic(nn.Module):
    """中心化Critic网络，输入全局状态"""
    def __init__(self, global_state_dim, num_agents, hidden_dim=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(global_state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, num_agents)  # 输出维度=智能体数量
        )
    
    def forward(self, global_state):
        return self.net(global_state)